package lookup

import (
	"image"
	"sync"
)

// 拆分给定的区间为num个区间
// begin 开始值
// end 结束值
// num 区间数量
// result 结果数组
// err 错误信息
func SplitRange(begin, end, num int) (result []int) {
	if begin > end {
		return []int{begin, end}
	}
	total := end - begin + 1
	if total <= num {
		return []int{begin, end}
	}
	perLoop := total / num
	remainder := total % num
	if remainder != 0 {
		perLoop++
	}
	result = make([]int, 0, num+1)
	result = append(result, begin)
	for i := 0; i < num; i++ {
		start := begin + i*perLoop
		end1 := start + perLoop
		if i < remainder {
			end1++
		}
		if end1 > end {
			end1 = end
		}
		result = append(result, end1)
	}
	return result
}

func lookupAllThead(imgBin *imageBinary, x1, y1, x2, y2 int, templateBin *imageBinary, m float64, theadNum int) ([]GPoint, error) {
	if theadNum <= 1 {
		return lookupAll(imgBin, x1, y1, x2, y2, templateBin, m)
	}
	templateWidth := templateBin.width
	templateHeight := templateBin.height

	result := SplitRange(x1, x2-templateWidth+1, theadNum)
	stepChan := make(chan *Step, theadNum)
	for i := 1; i < len(result); i++ {
		stepChan <- &Step{Begin: result[i-1], End: result[i]}
	}
	var list []GPoint
	resultPoint := make(chan *GPoint, theadNum)
	go func() {
		for g := range resultPoint {
			list = append(list, *g)
		}
	}()

	close(stepChan)
	wg := WaitGroupWrap(theadNum, stepChan, func(step *Step) {
		for x := step.Begin; x <= step.End; x++ {
			for y := y1; y <= y2-templateHeight+1; y++ {
				g, err := lookup(imgBin, templateBin, x, y, m)
				if err != nil {
					return
				}
				if g != nil {
					resultPoint <- g
				}
			}
		}
	})
	wg.Wait()
	close(resultPoint)
	return list, nil
}

type Step struct {
	Begin int
	End   int
}

func WaitGroupWrap[G any](maxThread int, chanAny chan *G, cb func(g *G)) *sync.WaitGroup {
	wg := new(sync.WaitGroup)
	wg.Add(maxThread)
	for i := 0; i < maxThread; i++ {
		go func() {
			for g := range chanAny {
				cb(g)
			}
			wg.Done()
		}()
	}
	return wg
}

// FindAllInRect searches for all occurrences of template only inside a part of the image.
// This can be used to speed up the search if you know the region of the
// image that the template should appear in.
func (l *Lookup) FindAllInRectThead(template image.Image, rect image.Rectangle, threshold float64, num int) ([]GPoint, error) {
	if len(l.imgBin.channels) == 1 {
		template = ensureGrayScale(template)
	}
	tb := newImageBinary(template)
	return lookupAllThead(l.imgBin, rect.Min.X, rect.Min.Y, rect.Max.X, rect.Max.Y, tb, threshold, num)
}

// FindAll searches for all occurrences of template inside the whole image.
func (l *Lookup) FindAllThead(template image.Image, threshold float64, num int) ([]GPoint, error) {
	return l.FindAllInRectThead(template, image.Rect(0, 0, l.imgBin.width-1, l.imgBin.height-1), threshold, num)
}
